package com.ase.jython;

import java.io.IOException;
import java.io.Reader;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;

import javax.script.AbstractScriptEngine;
import javax.script.Bindings;
import javax.script.Compilable;
import javax.script.CompiledScript;
import javax.script.Invocable;
import javax.script.ScriptContext;
import javax.script.ScriptEngine;
import javax.script.ScriptEngineFactory;
import javax.script.ScriptException;
import javax.script.SimpleBindings;

import org.python.core.CompileMode;
import org.python.core.Py;
import org.python.core.PyCode;
import org.python.core.PyObject;
import org.python.core.PySystemState;

public class JythonScriptEngine extends AbstractScriptEngine implements
		Compilable, Invocable {

	private PyObject myScope;
	private ScriptEngineFactory factory;
	public static final String JYTHON_COMPILE_MODE = "com.sun.script.jython.comp.mode";

	private static ThreadLocal<PySystemState> systemState;
	static {
		PySystemState.initialize();
		systemState = new ThreadLocal<PySystemState>();
	}

	public JythonScriptEngine() {
		myScope = newScope(context);
	}

	private PyObject newScope(ScriptContext ctx) {
		return new JythonObject(this, ctx);
	}

	public Object eval(String str, ScriptContext ctx) throws ScriptException {
		PyCode code = compileScript(str, ctx);
		return evalCode(code, ctx);
	}

	public Object eval(Reader reader, ScriptContext ctx) throws ScriptException {
		return eval(readFully(reader), ctx);
	}

	public Bindings createBindings() {
		return new SimpleBindings();
	}

	private class JythonCompiledScript extends CompiledScript {
		// my compiled code
		private PyCode code;

		JythonCompiledScript(PyCode code) {
			this.code = code;
		}

		public ScriptEngine getEngine() {
			return JythonScriptEngine.this;
		}

		public Object eval(ScriptContext ctx) throws ScriptException {
			return evalCode(code, ctx);
		}
	}

	// Compilable methods
	public CompiledScript compile(String script) throws ScriptException {
		PyCode code = compileScript(script, context);
		return new JythonCompiledScript(code);
	}

	public CompiledScript compile(Reader reader) throws ScriptException {
		return compile(readFully(reader));
	}

	// Invocable methods
	public Object invokeFunction(String name, Object... args)
			throws ScriptException, NoSuchMethodException {
		return invokeImpl(null, name, args);
	}

	public Object invokeMethod(Object obj, String name, Object... args)
			throws ScriptException, NoSuchMethodException {
		if (obj == null) {
			throw new IllegalArgumentException("script object is null");
		}
		return invokeImpl(obj, name, args);
	}

	private Object invokeImpl(Object obj, String name, Object... args)
			throws ScriptException, NoSuchMethodException {
		if (name == null) {
			throw new NullPointerException("method name is null");
		}
		setSystemState();

		PyObject thiz;
		if (obj instanceof PyObject) {
			thiz = (PyObject) obj;
		} else if (obj == null) {
			thiz = myScope;
		} else {
			thiz = java2py(obj);
		}

		PyObject func = thiz.__findattr__(name);
		if (func == null || !func.isCallable()) {
			if (thiz == myScope) {
				systemState.get();
				// lookup in built-in functions. This way
				// user can call invoke built-in functions.
				PyObject builtins = PySystemState.builtins;
				func = builtins.__finditem__(name);
			}
		}

		if (func == null || !func.isCallable()) {
			throw new NoSuchMethodException(name);
		}
		PyObject res = func.__call__(wrapArguments(args));
		return py2java(res);
	}

	public <T> T getInterface(Object obj, Class<T> clazz) {
		if (obj == null) {
			throw new IllegalArgumentException("script object is null");
		}
		return makeInterface(obj, clazz);
	}

	public <T> T getInterface(Class<T> clazz) {
		return makeInterface(null, clazz);
	}

	@SuppressWarnings("unchecked")
	private <T> T makeInterface(Object obj, Class<T> clazz) {
		if (clazz == null || !clazz.isInterface()) {
			throw new IllegalArgumentException("interface Class expected");
		}
		final Object thiz = obj;
		return (T) Proxy.newProxyInstance(clazz.getClassLoader(),
				new Class[] { clazz }, new InvocationHandler() {
					public Object invoke(Object proxy, Method m, Object[] args)
							throws Throwable {
						Object res = invokeImpl(thiz, m.getName(), args);
						return py2java(java2py(res), m.getReturnType());
					}
				});
	}

	// ScriptEngine methods
	public ScriptEngineFactory getFactory() {
		synchronized (this) {
			if (factory == null) {
				factory = new JythonScriptEngineFactory();
			}
		}
		return factory;
	}

	public void setContext(ScriptContext ctx) {
		super.setContext(ctx);
		// update myScope to keep it in-sync
		myScope = newScope(context);
	}

	// package-private methods
	void setFactory(ScriptEngineFactory factory) {
		this.factory = factory;
	}

	static PyObject java2py(Object javaObj) {
		return Py.java2py(javaObj);
	}

	static Object py2java(PyObject pyObj, @SuppressWarnings("rawtypes") Class type) {
		return (pyObj == null) ? null : pyObj.__tojava__(type);
	}

	static Object py2java(PyObject pyObj) {
		return py2java(pyObj, Object.class);
	}

	static PyObject[] wrapArguments(Object[] args) {
		if (args == null) {
			return new PyObject[0];
		}

		PyObject[] res = new PyObject[args.length];
		for (int i = 0; i < args.length; i++) {
			res[i] = java2py(args[i]);
		}
		return res;
	}

	// internals only below this point
	private PyObject getJythonScope(ScriptContext ctx) {
		if (ctx == context) {
			return myScope;
		} else {
			return newScope(ctx);
		}
	}

	private void setSystemState() {
		/*
		 * From my reading of Jython source, it appears that PySystemState is
		 * set on per-thread basis. So, I maintain it in a thread local and set
		 * it. Besides, this also helps in setting correct class loader -- which
		 * is thread context class loader.
		 */
		if (systemState.get() == null) {
			// we entering into this thread for the first time.

			PySystemState newState = new PySystemState();
			ClassLoader cl = Thread.currentThread().getContextClassLoader();
			newState.setClassLoader(cl);
			systemState.set(newState);
			Py.setSystemState(newState);
		}
	}

	private PyCode compileScript(String script, ScriptContext ctx)
			throws ScriptException {
		try {
			setSystemState();
			String fileName = (String) ctx.getAttribute(ScriptEngine.FILENAME);
			if (fileName == null) {
				fileName = "<unknown>";
			}

			/*
			 * Jython parser seems to have 3 input modes (called compile "kind")
			 * These are "single", "eval" and "exec". I don't clearly understand
			 * the difference. But, with "eval" and "exec" certain features are
			 * not working. For eg. with "eval" assignments are not working.
			 * I've used "exec". But, that is customizable by special attribute.
			 */
			String mode = (String) ctx.getAttribute(JYTHON_COMPILE_MODE);
			if (mode == null) {
				mode = "exec";
			}
			return (PyCode) Py.compile_flags(script, fileName,
					CompileMode.getMode(mode), Py.getCompilerFlags());

		} catch (Exception exp) {
			throw new ScriptException(exp);
		}
	}

	private Object evalCode(PyCode code, ScriptContext ctx)
			throws ScriptException {
		try {
			setSystemState();
			PyObject scope = getJythonScope(ctx);
			PyObject res = Py.runCode(code, scope, scope);
			return res.__tojava__(Object.class);
		} catch (Exception exp) {
			throw new ScriptException(exp);
		}
	}

	private String readFully(Reader reader) throws ScriptException {
		char[] arr = new char[8 * 1024]; // 8K at a time
		StringBuilder buf = new StringBuilder();
		int numChars;
		try {
			while ((numChars = reader.read(arr, 0, arr.length)) > 0) {
				buf.append(arr, 0, numChars);
			}
		} catch (IOException exp) {
			throw new ScriptException(exp);
		}
		return buf.toString();
	}

}
